import pymunk
import pymunk.util as util
from pymunk import Vec2d

import math
import sys
import random

class WorldArm:
    def __init__(self, world, attachedToBody):
        # Body1 is closest to the fixed point
        body1 = pymunk.Body(1, 1)
        body1.position = attachedToBody.position
        shape = pymunk.Circle(body1, 5)
        shape.elasticity = 0.95
        shape.group = 1
        world.space.add(body1, shape)
        world.balls.append(shape)
        world.arm1Shape1 = shape
    
        # Body2 is at the end of the arm
        body2 = pymunk.Body(1, 1)
        body2.position = attachedToBody.position + Vec2d(world.armLength, 0)
        shape = pymunk.Circle(body2, 5)
        shape.elasticity = 0.95
        shape.group = 1
        world.space.add(body2, shape)
        world.balls.append(shape)
        world.arm1Shape2 = shape
        
        # Constrain body1 to the fixed point
        rotationJoint = pymunk.PinJoint(body1, attachedToBody, (0,0), (0,0))    
        world.space.add(rotationJoint)
        
        # Constrain the two circles to move together
        upperJoint = pymunk.PinJoint(body2, body1, (0,5), (0,5))    
        lowerJoint = pymunk.PinJoint(body2, body1, (0,-5), (0,-5))    
        upperJoint2 = pymunk.PinJoint(body2, body1, (0,5), (0,-5))    
        lowerJoint2 = pymunk.PinJoint(body2, body1, (0,-5), (0,5))    
        world.space.add(upperJoint, lowerJoint, upperJoint2, lowerJoint2)
        
        world.visualBallConnections.append( (len(world.balls) - 1, len(world.balls) - 2) ) 
    
    